#!/usr/bin/env python

import time
from franky import Affine, CartesianMotion, Robot, ReferenceType
import numpy as np
from dataclasses import dataclass

@dataclass
class CartesianCmd:
    position: Affine
    max_speed: float


class FrankaPosServo:
    """
    A simple class to control Franka robot in Cartesian space.
    It can send commands to the robot in both blocking and non-blocking ways.
    """
    def __init__(self, ip: str = "172.16.0.2", dynamics_factor: float = 0.5, goal_tolerance: float = 0.001):
        
        self.robot_ip = ip
        self.robot = Robot(ip)
        self.robot.relative_dynamics_factor = dynamics_factor
        self.goal_tolerance = goal_tolerance
        
        print("Franka Servo initialized successfully!")
    
    def __del__(self):
        """
        Destructor to stop the robot when the object is deleted.
        """
        self.stop_robot()
        print("Franka Servo stopped successfully!")
    
    def gen_abs_move_cmd(self, goal_position, max_speed: float = 0.01):
        """        
        Generate a command to move the robot to an absolute position in Cartesian space.
        Args:
        ---
            goal_position: The target position in Cartesian space, should be a 3d array-like.
            max_speed: The maximum speed of the robot, default is 0.01 m/s.
        Returns:
        ---
            A CartesianCmd object that can be published to the robot.
        """

        goal_position, max_speed = self._type_check(goal_position, max_speed)
        return CartesianCmd(
            position=Affine(goal_position, quaternion=[1, 0, 0, 0]),
            max_speed=max_speed
        )
    
    def gen_rela_move_cmd(self, relative_move, max_speed: float = 0.01):
        """
        Generate a command to move the robot by a relative offset in Cartesian space.
        Args:
            relative_move: The relative offset in Cartesian space, should be a 3d array-like.
            max_speed: The maximum speed of the robot, default is 0.01 m/s.
        Returns:
            A CartesianCmd object that can be published to the robot.
        """

        relative_move, max_speed = self._type_check(relative_move, max_speed)
        ee_pose = self.robot.current_cartesian_state.pose.end_effector_pose
        goal_position = Affine(relative_move) * ee_pose
        return CartesianCmd(
            position=goal_position,
            max_speed=max_speed
        )
    
    def non_blocking_move(self, cmd: CartesianCmd):
        """
        Send a command to the robot in a non-blocking way.
        The function will return immediately after sending the command.
        Args:
            cmd: A CartesianCmd object that can be published to the robot.
        """
        try:
            self.robot.translation_velocity_limit.set(cmd.max_speed)
        except:
            print("Failed to set translation velocity limit, using current value.")
        self.robot.move(CartesianMotion(cmd.position, ReferenceType.Absolute), asynchronous=True)
        return cmd.position.translation  
        
    
    def blocking_move(self, cmd: CartesianCmd):
        """
        Send a command to the robot in a blocking way.
        The function will block until the command is finished.
        Args:
            cmd: A CartesianCmd object that can be published to the robot.
        """
        try:
            self.robot.translation_velocity_limit.set(cmd.max_speed)
        except:
            print("Failed to set translation velocity limit, using current value.")
        self.robot.move(CartesianMotion(cmd.position, ReferenceType.Absolute), asynchronous=False)
    
    def is_cmd_finished(self, goal_position: np.ndarray) -> bool:
        """
        Check if the robot has reached the goal position.
        Args:
            goal_position: The target position in Cartesian space, should be a 3d array-like
        Returns:
            True if the robot has reached the goal position, False otherwise.
        """
        now_position = self.robot_position
        if np.linalg.norm(now_position - goal_position) < self.goal_tolerance:
            return True
        vel = self.robot.current_cartesian_state.velocity.end_effector_twist.linear
        if np.linalg.norm(now_position - goal_position) < self.goal_tolerance * 1.5 and np.linalg.norm(vel) < 0.001:
            return True
        return False
    
    def stop_robot(self):
        """
        Stop the robot immediately.
        This will stop any ongoing motion and set the robot to a safe state.
        """
        self.robot.stop()
    
    def wait_motion_finish(self):
        """
        Wait for the robot to finish its current motion.
        This will block until the robot has stopped moving.
        If the robot encounters an error during motion, and it is moving in an unblocking way,
        an exception will not be raised until this function is called.
        Therefore, it is recommended to call this function after sending a non-blocking command
        """
        self.robot.join_motion()

    
    @property
    def robot_position(self):
        """
        Get the current position of the robot in Cartesian space.
        Returns:
            The current position of the robot as an ndarray.
        """
        cartesian_state = self.robot.current_cartesian_state
        return cartesian_state.pose.end_effector_pose.translation
    
    @property
    def robot_pose(self):
        """
        Get the current pose of the robot in Cartesian space.
        Returns:
            The current pose of the robot as an Affine object.
        """
        cartesian_state = self.robot.current_cartesian_state
        return cartesian_state.pose.end_effector_pose

    
    def _type_check(self, goal_position, max_speed):
        if type(goal_position) == list:
            assert len(goal_position) == 3, "goal_position should be a 3d array-like."
            goal_position = np.array(goal_position)
        elif type(goal_position) == np.ndarray:
            assert goal_position.size == 3, "goal_position should be a 3d array-like."
            goal_position.squeeze()
        else:
            raise TypeError("goal position should be a list or a ndarray.")
        
        if max_speed <= 0:
            raise ValueError("max_speed should be a positive number.")
    
        return goal_position, max_speed
            

def main():
    """
    Main function for demonstration purposes.
    This function can be called from command line after installation.
    """
    print("Franka Servo Demo")
    print("Initializing Franka robot...")
    
    try:
        franka_servo = FrankaPosServo()
        original_position = franka_servo.robot_position
        print(f"Robot current position: {original_position}")
        
        print("Moving robot up by 5cm...")
        cmd = franka_servo.gen_rela_move_cmd([0, 0, 0.05])
        franka_servo.non_blocking_move(cmd)
        time.sleep(3)
        
        print("Moving robot back to original position...")
        new_cmd = franka_servo.gen_abs_move_cmd(original_position)
        franka_servo.blocking_move(new_cmd)
        
        print("Demo completed successfully!")
        
    except Exception as e:
        print(f"Error during demo: {e}")
        return 1
    
    return 0


if __name__ == "__main__":
    import sys
    sys.exit(main())

